package com.ruoyi.teaching.flink;

import com.ruoyi.common.core.utils.SpringUtils;
import com.ruoyi.teaching.domain.YaHighInterview;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;

/**
 * MySQL数据源
 */
@Slf4j
public class MySQLSource extends RichSourceFunction<YaHighInterview> {
    private volatile boolean isRunning = true;
    private transient DataSource dataSource;
    public static final int BATCH_SIZE = 100;
    private static final String QUERY_DATA = "SELECT * FROM ya_high_interview";
    private static final String QUERY_COUNT = "SELECT count(id) FROM ya_high_interview";
    //保证线程之间的可见性
    @Getter
    private static volatile int percentage = 0;

    @Override
    public void open(Configuration parameters) {
        // 初始化数据源
        dataSource = SpringUtils.getBean(DataSource.class);
    }

    @Override
    public void run(SourceContext<YaHighInterview> ctx) throws Exception {
        // 从MySQL读取数据并发送到Flink流处理
        log.debug("开始从MySQL读取数据...");
        percentage = 0;
        try (Connection conn = dataSource.getConnection()) {
            //先查询总记录数
            PreparedStatement stmtCount = conn.prepareStatement(QUERY_COUNT,
                    ResultSet.TYPE_FORWARD_ONLY,
                    ResultSet.CONCUR_READ_ONLY);
            ResultSet rsCount = stmtCount.executeQuery();
            if (!rsCount.next()) {
                return;
            }
            double total = rsCount.getInt(1) * 0.01;

            PreparedStatement stmt = conn.prepareStatement(QUERY_DATA,
                    ResultSet.TYPE_FORWARD_ONLY,
                    ResultSet.CONCUR_READ_ONLY);
            stmt.setFetchSize(BATCH_SIZE); // 设置每次获取的记录数
            ResultSet rs = stmt.executeQuery(); // 执行查询
            int count = 0;
            // 循环读取每一行数据
            while (isRunning && rs.next()) {
                YaHighInterview yaHighInterview = new YaHighInterview();
                yaHighInterview.setQuestion(rs.getString("question")); // 设置问题
                yaHighInterview.setAnswer(rs.getString("answer")); // 设置答案
                yaHighInterview.setKnowledgePoint(rs.getString("knowledge_point")); // 设置知识点
                yaHighInterview.setImportance(rs.getLong("importance")); // 设置重要性
                yaHighInterview.setStage(rs.getLong("stage")); // 设置阶段
                yaHighInterview.setId(rs.getLong("id")); // 设置ID
                ctx.collect(yaHighInterview); // 发送到Flink流处理
                count++;
                // 每读取BATCH_SIZE条记录打印一次日志
                if (count % BATCH_SIZE == 0) {
                    log.debug("已读取 {} 条记录", count);
                }
                //进度=count/total
                percentage = (int) (count / (total));
            }
            log.debug("MySQL数据读取完成，总共读取 {} 条记录", count);
        }
    }

    @Override
    public void cancel() {
        // 取消数据读取
        isRunning = false;
    }

    @Override
    public void close() {
        // 关闭资源
    }
}
